package datastore

import (
	"context"

	"github.com/pkg/errors"
	namespaceDS "github.com/stackrox/rox/central/namespace/datastore"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/cache"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/common"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store"
	"github.com/stackrox/rox/central/vulnmgmt/vulnerabilityrequest/utils"
	v1 "github.com/stackrox/rox/generated/api/v1"
	"github.com/stackrox/rox/generated/storage"
	"github.com/stackrox/rox/pkg/errox"
	"github.com/stackrox/rox/pkg/grpc/authn"
	"github.com/stackrox/rox/pkg/protocompat"
	"github.com/stackrox/rox/pkg/sac"
	"github.com/stackrox/rox/pkg/sac/resources"
	"github.com/stackrox/rox/pkg/search"
)

var (
	requesterOrApproverSAC = sac.ForResources(sac.ForResource(resources.VulnerabilityManagementRequests), sac.ForResource(resources.VulnerabilityManagementApprovals))
	approverSAC            = sac.ForResource(resources.VulnerabilityManagementApprovals)
	requesterSAC           = sac.ForResource(resources.VulnerabilityManagementRequests)
	allAccessCtx           = sac.WithAllAccess(context.Background())
)

type datastoreImpl struct {
	pendingReqCache    cache.VulnReqCache
	activeReqCache     cache.VulnReqCache
	store              store.Store
	namespaceDatastore namespaceDS.DataStore
}

func (ds *datastoreImpl) Search(ctx context.Context, q *v1.Query) ([]search.Result, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)

	return ds.store.Search(storeCtx, q)
}

func (ds *datastoreImpl) SearchRequests(ctx context.Context, q *v1.Query) ([]*v1.SearchResult, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)

	// TODO(ROX-29943): remove 2 pass database queries
	results, err := ds.Search(storeCtx, q)
	if err != nil {
		return nil, err
	}

	requests, missingIndices, err := ds.store.GetMany(storeCtx, search.ResultsToIDs(results))
	if err != nil {
		return nil, err
	}
	results = search.RemoveMissingResults(results, missingIndices)
	return convertMany(requests, results)
}

func (ds *datastoreImpl) SearchRawRequests(ctx context.Context, q *v1.Query) ([]*storage.VulnerabilityRequest, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)

	var requests []*storage.VulnerabilityRequest
	err := ds.store.GetByQueryFn(storeCtx, q, func(request *storage.VulnerabilityRequest) error {
		requests = append(requests, request)
		return nil
	})
	if err != nil {
		return nil, err
	}

	return requests, nil
}

func (ds *datastoreImpl) Count(ctx context.Context, q *v1.Query) (int, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return 0, err
	} else if !ok {
		return 0, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)

	return ds.store.Count(storeCtx, q)
}

func (ds *datastoreImpl) Exists(ctx context.Context, id string) (bool, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return false, err
	} else if !ok {
		return false, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)
	return ds.store.Exists(storeCtx, id)
}

func (ds *datastoreImpl) Get(ctx context.Context, id string) (*storage.VulnerabilityRequest, bool, error) {
	if id == "" {
		return nil, false, errors.New("vulnerability request ID must be provided")
	}
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return nil, false, err
	} else if !ok {
		return nil, false, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)
	return ds.store.Get(storeCtx, id)
}

func (ds *datastoreImpl) GetMany(ctx context.Context, ids []string) ([]*storage.VulnerabilityRequest, error) {
	if ok, err := requesterOrApproverSAC.ReadAllowedToAny(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)
	ret, _, err := ds.store.GetMany(storeCtx, ids)
	return ret, err
}

func (ds *datastoreImpl) AddRequest(ctx context.Context, request *storage.VulnerabilityRequest) error {
	// Only those with WRITE on VulnerabilityManagementRequests are allowed to add a _new_ request
	if ok, err := requesterSAC.WriteAllowed(ctx); err != nil {
		return err
	} else if !ok {
		return sac.ErrResourceAccessDenied
	}

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	return ds.store.Upsert(storeReadWriteCtx, request)
}

func (ds *datastoreImpl) UpdateRequestStatus(ctx context.Context, id, message string, status storage.RequestStatus) (*storage.VulnerabilityRequest, error) {
	if id == "" {
		return nil, errors.New("vulnerability request ID must be provided")
	}
	if status == storage.RequestStatus_PENDING || status == storage.RequestStatus_APPROVED_PENDING_UPDATE {
		return nil, errors.New("vulnerability request cannot be moved to pending state")
	}
	if message == "" {
		return nil, errors.New("comment is required when approving/denying a vulnerability request")
	}

	// Only those with WRITE on VulnerabilityManagementApprovals are allowed to update status since this is an approval or denial flow
	if ok, err := approverSAC.WriteAllowed(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	req, found, err := ds.store.Get(storeReadWriteCtx, id)
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, errors.Errorf("vulnerability request %s not found", id)
	}

	if status == storage.RequestStatus_APPROVED && utils.V2GlobalScope(req.GetScope()) {
		err = ds.ensureApproverCanSeeAllImages(ctx)
		if err != nil {
			return nil, err
		}
	}

	req.LastUpdated = protocompat.TimestampNow()
	approver := authn.UserFromContext(ctx)
	req.Approvers = append(req.Approvers, approver)
	req.ApproversV2 = append(req.ApproversV2, utils.ApproverV2(approver))
	req.Comments = append(req.Comments, utils.CreateRequestCommentProto(ctx, message))

	// If this was an update whose status is being updated, the request fields must be changed appropriately.
	if req.GetStatus() == storage.RequestStatus_APPROVED_PENDING_UPDATE {
		// `APPROVED_PENDING_UPDATE` indicates that the original request was approved and its update was pending.
		// Set the status of the original request back to APPROVED whether the update was approved or denied.
		req.Status = storage.RequestStatus_APPROVED

		// If the update was approved, apply the update to original request and wipe out the update field.
		if status == storage.RequestStatus_APPROVED {
			applyUpdateRequest(req)
		}

		// Wipe out the update field whether the update was approved or denied.
		req.UpdatedReq = nil
	} else {
		req.Status = status
	}

	// If the request is denied, mark it as expired (due to denial) so that it is not mistakenly pulled out when searching by expiry.
	if req.GetStatus() == storage.RequestStatus_DENIED {
		req.Expired = true
	}

	if err := ds.store.Upsert(storeReadWriteCtx, req); err != nil {
		return nil, err
	}
	return req, nil
}

func (ds *datastoreImpl) ensureApproverCanSeeAllImages(ctx context.Context) error {
	imageSAC := sac.ForResource(resources.Image)
	hasReadAll, err := imageSAC.HasGlobalRead(ctx)
	if err != nil {
		return err
	}
	if hasReadAll {
		return nil
	}
	allNamespaces, err := ds.namespaceDatastore.GetAllNamespaces(allAccessCtx)
	if err != nil {
		return err
	}
	filteredNamespaces, err := imageSAC.FilterAccessibleNamespaces(ctx, storage.Access_READ_ACCESS, allNamespaces)
	if err != nil {
		return err
	}
	if len(allNamespaces) == len(filteredNamespaces) {
		return nil
	}
	return errors.Wrap(errox.NotAuthorized, "Access scope of global exception request approvers should include all namespaces")
}

func applyUpdateRequest(req *storage.VulnerabilityRequest) {
	switch r := req.GetUpdatedReq().(type) {
	case *storage.VulnerabilityRequest_DeferralUpdate:
		req.Req = &storage.VulnerabilityRequest_DeferralReq{
			DeferralReq: &storage.DeferralRequest{
				Expiry: r.DeferralUpdate.GetExpiry(),
			},
		}
		req.Entities = &storage.VulnerabilityRequest_Cves{
			Cves: &storage.VulnerabilityRequest_CVEs{
				Cves: r.DeferralUpdate.GetCVEs(),
			},
		}
	case *storage.VulnerabilityRequest_FalsePositiveUpdate:
		req.Entities = &storage.VulnerabilityRequest_Cves{
			Cves: &storage.VulnerabilityRequest_CVEs{
				Cves: r.FalsePositiveUpdate.GetCVEs(),
			},
		}
	}
}

func (ds *datastoreImpl) UpdateRequestExpiry(ctx context.Context, id, message string, updatedExpiry *storage.RequestExpiry) (*storage.VulnerabilityRequest, error) {
	if id == "" {
		return nil, errors.New("vulnerability request ID must be provided")
	}
	if message == "" {
		return nil, errors.New("comment is required when updating a request's expiry")
	}

	req, err := ds.ensureApproverOrInitialRequestorReturningRequest(ctx, id)
	if err != nil {
		return nil, err
	}

	if req.GetTargetState() != storage.VulnerabilityState_DEFERRED {
		return nil, errors.Errorf("request %s is not a deferral request thus doesn't have an expiry to update", req.GetId())
	}

	if req.GetExpired() || req.GetStatus() == storage.RequestStatus_DENIED {
		return nil, errors.Errorf("request %s already expired or denied thus cannot be updated", req.GetId())
	}

	req.LastUpdated = protocompat.TimestampNow()
	req.Comments = append(req.Comments, utils.CreateRequestCommentProto(ctx, message))

	switch req.GetStatus() {
	// If request was already pending, then just overwrite the request
	case storage.RequestStatus_PENDING:
		req.GetDeferralReq().Expiry = updatedExpiry

	// If an update request was already pending, then just overwrite the update request
	case storage.RequestStatus_APPROVED_PENDING_UPDATE:
		req.GetUpdatedDeferralReq().Expiry = updatedExpiry

	// If request was already approved, then create a new update request
	case storage.RequestStatus_APPROVED:
		updatedDeferral := req.GetDeferralReq().CloneVT() // everything but the expiry is the same, so just clone the original
		updatedDeferral.Expiry = updatedExpiry
		req.UpdatedReq = &storage.VulnerabilityRequest_UpdatedDeferralReq{UpdatedDeferralReq: updatedDeferral}
		req.Status = storage.RequestStatus_APPROVED_PENDING_UPDATE
	}

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	if err := ds.store.Upsert(storeReadWriteCtx, req); err != nil {
		return nil, err
	}
	return req, nil
}

func (ds *datastoreImpl) UpdateRequest(ctx context.Context, id string, update *common.UpdateRequest) (*storage.VulnerabilityRequest, error) {
	req, err := ds.ensureApproverOrInitialRequestorReturningRequest(ctx, id)
	if err != nil {
		return nil, err
	}

	if req.GetExpired() || req.GetStatus() == storage.RequestStatus_DENIED {
		return nil, errors.Errorf("request %s cannot be updated because it has expired or has been denied", req.GetId())
	}

	if (req.GetTargetState() == storage.VulnerabilityState_DEFERRED && update.FalsePositiveUpdate != nil) ||
		(req.GetTargetState() == storage.VulnerabilityState_FALSE_POSITIVE && update.DeferralUpdate != nil) {
		return nil, errors.Errorf("exception %s cannot be updated. Exception type(defer/false positive) cannot be changed", req.GetId())
	}

	if utils.IsUpdateNoOp(req, update) {
		return nil, errors.Errorf("the update does not change the exception %s", req.GetId())
	}

	req.LastUpdated = protocompat.TimestampNow()
	req.Comments = append(req.Comments, utils.CreateRequestCommentProto(ctx, update.Comment))
	if update.DeferralUpdate != nil {
		deferralUpdate(req, update.DeferralUpdate)
	} else if update.FalsePositiveUpdate != nil {
		falsePositiveUpdate(req, update.FalsePositiveUpdate)
	}

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	if err := ds.store.Upsert(storeReadWriteCtx, req); err != nil {
		return nil, err
	}
	return req, nil
}

func (ds *datastoreImpl) ensureApproverOrInitialRequestorReturningRequest(ctx context.Context, id string) (*storage.VulnerabilityRequest, error) {
	// Only those with WRITE on either VulnerabilityManagementApprovals or VulnerabilityManagementRequests can make this change
	if ok, err := requesterOrApproverSAC.WriteAllowedToAny(ctx); err != nil {
		return nil, err
	} else if !ok {
		return nil, sac.ErrResourceAccessDenied
	}

	storeReadCtx := getStoreAccessContext(ctx, storage.Access_READ_ACCESS)
	req, found, err := ds.store.Get(storeReadCtx, id)
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, errors.Errorf("vulnerability request %s not found", id)
	}

	// If the user has VulnerabilityManagementRequests but not VulnerabilityManagementApprovals, then this can only be done if current user is the requester
	if ok, err := requesterSAC.WriteAllowed(ctx); err == nil && ok {
		if ok, err := approverSAC.WriteAllowed(ctx); err != nil || !ok {
			user, err := authn.IdentityFromContext(ctx)
			if err != nil {
				return nil, err
			}
			if req.GetRequestor().GetId() != user.UID() {
				return nil, errors.Wrap(sac.ErrResourceAccessDenied, "requests can only be updated by the original requester")
			}
		}
	}

	return req, nil
}

func deferralUpdate(req *storage.VulnerabilityRequest, update *storage.DeferralUpdate) {
	switch req.GetStatus() {
	// If request was already pending, then just overwrite the request
	case storage.RequestStatus_PENDING:
		req.GetDeferralReq().Expiry = update.GetExpiry()
		req.GetCves().Cves = update.GetCVEs()

	// If an update request was already pending, then just overwrite the update request
	case storage.RequestStatus_APPROVED_PENDING_UPDATE:
		req.UpdatedReq = &storage.VulnerabilityRequest_DeferralUpdate{DeferralUpdate: update}

	// If request was already approved, then create a new update request
	case storage.RequestStatus_APPROVED:
		req.UpdatedReq = &storage.VulnerabilityRequest_DeferralUpdate{DeferralUpdate: update}
		req.Status = storage.RequestStatus_APPROVED_PENDING_UPDATE
	}
}

func falsePositiveUpdate(req *storage.VulnerabilityRequest, update *storage.FalsePositiveUpdate) {
	switch req.GetStatus() {
	// If request was already pending, then just overwrite the request
	case storage.RequestStatus_PENDING:
		req.GetCves().Cves = update.GetCVEs()

	// If an update request was already pending, then just overwrite the update request
	case storage.RequestStatus_APPROVED_PENDING_UPDATE:
		req.UpdatedReq = &storage.VulnerabilityRequest_FalsePositiveUpdate{FalsePositiveUpdate: update}

	// If request was already approved, then create a new update request
	case storage.RequestStatus_APPROVED:
		req.UpdatedReq = &storage.VulnerabilityRequest_FalsePositiveUpdate{FalsePositiveUpdate: update}
		req.Status = storage.RequestStatus_APPROVED_PENDING_UPDATE
	}
}

func (ds *datastoreImpl) MarkRequestInactive(ctx context.Context, id, comment string) (*storage.VulnerabilityRequest, error) {
	// Only those with WRITE on either VulnerabilityManagementApprovals or VulnerabilityManagementRequests can make this change
	req, err := ds.ensureApproverOrInitialRequestorReturningRequest(ctx, id)
	if err != nil {
		return nil, err
	}

	if req.GetExpired() {
		if req.GetStatus() == storage.RequestStatus_DENIED {
			return nil, errors.Errorf("vulnerability request %s has expired due to request denial. Undo is noop", id)
		}
		return nil, errors.Errorf("vulnerability request %s has expired. Undo is noop", id)
	}

	req.LastUpdated = protocompat.TimestampNow()
	req.Comments = append(req.Comments, utils.CreateRequestCommentProto(ctx, comment))
	req.Expired = true

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	if err := ds.store.Upsert(storeReadWriteCtx, req); err != nil {
		return nil, err
	}
	return req, nil
}

func (ds *datastoreImpl) RemoveRequest(ctx context.Context, id string) error {
	// Only those with WRITE VulnerabilityManagementRequests can cancel an existing pending request
	if ok, err := requesterSAC.WriteAllowed(ctx); err != nil {
		return err
	} else if !ok {
		return sac.ErrResourceAccessDenied
	}

	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	req, found, err := ds.store.Get(storeReadWriteCtx, id)
	if err != nil {
		return err
	}
	if !found {
		return errors.Errorf("vulnerability request %s not found", id)
	}

	// Requests can only be cancelled by the original requester
	user, err := authn.IdentityFromContext(ctx)
	if err != nil {
		return err
	}
	if req.GetRequestor().GetId() != user.UID() {
		return errors.Wrap(sac.ErrResourceAccessDenied, "requests can only be cancelled by the original requester")
	}

	// Only pending exceptions and pending updates to an existing exception can be deleted through API. All other
	// exceptions can only is cancelled and are retained in the system according to the retention configuration.
	switch req.GetStatus() {
	case storage.RequestStatus_PENDING:
		if err := ds.store.Delete(storeReadWriteCtx, id); err != nil {
			return errors.Wrapf(err, "could not delete vulnerability exception %s", id)
		}

	case storage.RequestStatus_APPROVED_PENDING_UPDATE:
		// Removing an update should take it back to its previous state. Since the original request object maintains
		// the original state, it can't be simply removed.
		req.Status = storage.RequestStatus_APPROVED
		req.UpdatedReq = nil
		if err := ds.store.Upsert(storeReadWriteCtx, req); err != nil {
			return errors.Wrapf(err, "could not delete vulnerability exception %s update", id)
		}
	default:
		return errors.Errorf("cannot delete exception %s. Only vulnerability exceptions or updates in pending state can be deleted", id)
	}
	return nil
}

// FOR INTERNAL USE ONLY
func (ds *datastoreImpl) RemoveRequestsInternal(ctx context.Context, ids []string) error {
	if ok, err := requesterSAC.WriteAllowed(ctx); err != nil {
		return err
	} else if !ok {
		return sac.ErrResourceAccessDenied
	}
	storeReadWriteCtx := getStoreAccessContext(ctx, storage.Access_READ_WRITE_ACCESS)
	if err := ds.store.DeleteMany(storeReadWriteCtx, ids); err != nil {
		return err
	}
	ds.pendingReqCache.RemoveMany(ids...)
	ds.activeReqCache.RemoveMany(ids...)
	return nil
}

// region helpers

func getStoreAccessContext(ctx context.Context, mode storage.Access) context.Context {
	accessModes := []storage.Access{mode}
	if mode == storage.Access_READ_WRITE_ACCESS {
		accessModes = append(accessModes, storage.Access_READ_ACCESS)
	}
	return sac.WithGlobalAccessScopeChecker(
		ctx,
		sac.AllowFixedScopes(
			sac.AccessModeScopeKeys(accessModes...),
			sac.ResourceScopeKeys(resources.VulnerabilityRequest),
		),
	)
}

func convertMany(vulnRequests []*storage.VulnerabilityRequest, results []search.Result) ([]*v1.SearchResult, error) {
	if len(vulnRequests) != len(results) {
		return nil, errors.Errorf("expected %d vuln requests but got %d", len(results), len(vulnRequests))
	}

	outputResults := make([]*v1.SearchResult, len(vulnRequests))
	for index, sar := range vulnRequests {
		outputResults[index] = convertOne(sar, &results[index])
	}
	return outputResults, nil
}

func convertOne(vulnRequest *storage.VulnerabilityRequest, result *search.Result) *v1.SearchResult {
	return &v1.SearchResult{
		Category:       v1.SearchCategory_VULN_REQUEST,
		Id:             vulnRequest.GetId(),
		Name:           vulnRequest.GetId(), // Requests do not have names.
		FieldToMatches: search.GetProtoMatchesMap(result.Matches),
		Score:          result.Score,
	}
}

// endregion helpers
